import logging
from pinecone import Pinecone
from sklearn.datasets import load_digits
import numpy as np
from collections import Counter
from tqdm import tqdm
import datetime

# 配置logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s')

# 连接到 Pinecone
api_key = "d86bcb40-a767-490b-b198-d5101fefbe7b"
pinecone = Pinecone(api_key=api_key)

# 索引名称
index_name = "mnist-index"

# 连接到索引
index = pinecone.Index(index_name)

# 从 scikit-learn 库中导入 MNIST 数据集
digits = load_digits(n_class=10)
X = digits.data
y = digits.target

# 分割数据集为训练集和测试集
split_idx = int(len(X) * 0.8)
X_train, X_test = X[:split_idx], X[split_idx:]
y_train, y_test = y[:split_idx], y[split_idx:]

# 准备训练数据
vectors = []
for i in tqdm(range(len(X_train)), desc='Preparing training data'):
    vector_id = str(i)
    vector_values = X_train[i].tolist()
    metadata = {"label": int(y_train[i])}
    vectors.append((vector_id, vector_values, metadata))

# 批量上传数据到 Pinecone
batch_size = 1000
num_batches = (len(vectors) + batch_size - 1) // batch_size
for i in tqdm(range(0, len(vectors), batch_size), desc='Uploading data', total=num_batches):
    batch = vectors[i:i + batch_size]
    index.upsert(batch)
logging.info(f"Successfully created index and uploaded {len(vectors)} pieces of data.")

# 准备查询数据并计算准确率
correct_predictions = 0
total_predictions = 0
for i in tqdm(range(len(X_test)), desc='Testing accuracy', total=len(X_test)):
    query_data = X_test[i].tolist()
    results = index.query(vector=query_data, top_k=11, include_metadata=True)
    labels = [match['metadata']['label'] for match in results['matches'] if 'metadata' in match and 'label' in match['metadata']]
    if y_test[i] in labels:
        correct_predictions += 1
    total_predictions += 1

accuracy = correct_predictions / total_predictions
logging.info(f"With k=11, the accuracy using Pinecone is {accuracy:.2f}.")

# 打印日志信息
logging.info(f"Successfully created index and uploaded {len(vectors)} pieces of data.")
logging.info(f"With k=11, the accuracy using Pinecone is {accuracy:.2f}.")